from dotenv import load_dotenv
from fastapi import FastAPI
from langchain_core.messages import SystemMessage, HumanMessage

from langchain_community.llms.tongyi import Tongyi
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langserve import add_routes


# 加载环境变量
load_dotenv()

# 1.创建模型
model = Tongyi()

# 2.准备 prompt
msg = [
    SystemMessage(content="请将以下的内容翻译为英语。"),
    HumanMessage(content='你好，请问你要去哪里？')
]

# result = model.invoke(msg)
# print(result)   # 返回content和response_metadata信息


# 解析返回的数据 [可以解析为字符串、JSON、自定义的POJO等]
# 3.创建返回的数据解析器
parser = StrOutputParser()
# print(parser.invoke(result))    # 只返回content的内容了

# 使用提示词模板
prompt_template = ChatPromptTemplate.from_messages([
    ('system', '请将以下的内容翻译为{language}'),
    ('user', '{text}'),
])

# 4.得到链
# chain = model | parser
chain = prompt_template | model | parser

# 5.直接用chain来调用
# print(chain.invoke(msg))    # 直接返回简洁的内容
print(chain.invoke({'language': 'English', 'text': '晚饭吃什么？川菜还是本帮菜？'}))  # 直接返回简洁的内容


# 把程序部署为服务
# 创建FastApi的应用
app = FastAPI(title='我的LangChain服务', version='v1.0.0', description='使用LangChain翻译任何语句的服务')

add_routes(
    app,
    chain,
    path='/chainDemo',
)

if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, host='localhost', port=8000)

# 请求路径 尾部invoke为固定值
# http://localhost:8000/chainDemo/invoke
# 请求参数
# {
#     "input": {
#         "language": "English",
#         "text": "晚饭吃什么？川菜还是本帮菜？"
#     }
# }
